package com.bigdata.yarn;

import org.apache.hadoop.conf.Configuration;
import org.apache.hadoop.security.UserGroupInformation;
import org.apache.hadoop.util.ToolRunner;
import org.apache.hadoop.yarn.api.records.ApplicationReport;
import org.apache.hadoop.yarn.api.records.QueueInfo;
import org.apache.hadoop.yarn.client.api.YarnClient;
import org.apache.hadoop.yarn.client.cli.RMAdminCLI;

import java.util.Date;
import java.util.List;
import java.util.Set;


/**
 * @Description: admin管理Yarn的管理器
 * @author: Shinelon
 * @date: 2017/10/17 14:27
 * @e-mail: 13971546488@163.com
 */

public class YarnAdminManager {

    private Configuration configuration;
    private RMAdminCLI rmAdminCLI;
    private YarnClient yarnClient;

    private static final String PARAM_REFRESH_QUEUES = "-refreshQueues";
    
    /** keytab文件 **/
	private static final String KEYTAB_FILE = "/Kerberos/hdfs.keytab";
	/** 主体名，与keytab文件中的主体保持一致 **/
	private static final String LOGIN_PRINCIPAL = "hdfs@HADOOP.COM";
	/** krb5.conf配置文件 **/
	private static final String KRB5_CONF = "/Kerberos/krb5.conf";

    public YarnAdminManager(Configuration configuration) {
        this.configuration = configuration;
        this.rmAdminCLI = new RMAdminCLI();
        this.yarnClient = YarnClient.createYarnClient();
        yarnClient.init(configuration);
        yarnClient.start();
        Configuration conf = yarnClient.getConfig();
        System.out.println(conf);
        System.out.println(configuration.get("yarn.resourcemanager.principal"));
    }

    public int refreshQueues() {
        String[] paramArr = {PARAM_REFRESH_QUEUES};
        int rtnCode = -1;
        try{
            rtnCode = ToolRunner.run(configuration, rmAdminCLI, paramArr);
        } catch(Exception e){
            e.printStackTrace();
        }
        return rtnCode;
    }

    public void printReport() throws  Exception{
        System.out.println("\n==================================================================");
        System.out.println(yarnClient);
        List<QueueInfo> queueList = yarnClient.getAllQueues();
        for (QueueInfo info : queueList) {
            System.out.println(info);
        }
        List<ApplicationReport> reportList = yarnClient.getApplications();
        for (ApplicationReport report : reportList) {
            String appId = report.getApplicationId().toString();
            String appName = report.getName();
            String state = report.getYarnApplicationState().name();
            String progress = report.getProgress() + "";
            String startTime = new Date(report.getStartTime()) + "";
            String appType = report.getApplicationType();
            String finishTime = new Date(report.getFinishTime()) + "";
            String host = report.getHost();
            String user = report.getUser();
            Set<String> tagSet = report.getApplicationTags();
            String tagStr = "";
            for (String tag : tagSet) {
                tagStr += tag;
            }
            System.out.println("appId:" + appId + ",appName:" + appName + ",state:" + state
                    + ",progress:" + progress + ",startTime:" + startTime + ",finishTime:" + finishTime
                    + ",appType:" + appType + ",host:" + host + ",user:" + user + ",tagStr:" + tagStr);
        }
        System.out.println("==================================================================\n");
    }
    
    public static void main(String[] args) throws Exception {
    	System.setProperty("java.security.krb5.conf", KRB5_CONF);
    	Configuration conf = new Configuration();
    	conf.set("hadoop.security.authentication", "kerberos");
		UserGroupInformation.setConfiguration(conf);
		UserGroupInformation.loginUserFromKeytab(LOGIN_PRINCIPAL,KEYTAB_FILE);
		YarnAdminManager yam = new YarnAdminManager(conf);
    	yam.printReport();
	}
}
